package main

import (
	"bytes"
	"flag"
	"fmt"
	"golang.org/x/image/draw"
	"image"
	"image/jpeg"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	// 定义命令行参数
	inputDir := flag.String("input", "input", "输入的文件夹路径")
	outputDir := flag.String("output", "output", "输出的文件夹路径")
	maxSize := flag.Int("maxSize", 32*1024, "压缩后的图片最大大小（字节）")
	flag.Parse()

	// 确保输出目录存在
	if err := os.MkdirAll(*outputDir, os.ModePerm); err != nil {
		log.Fatalf("无法创建输出目录: %v", err)
	}

	// 遍历输入目录中的所有文件和子目录
	err := filepath.Walk(*inputDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		// 检查文件扩展名是否为图片
		ext := strings.ToLower(filepath.Ext(path))
		if ext != ".jpg" && ext != ".jpeg" && ext != ".png" {
			return nil
		}

		// 读取图片文件
		imgData, err := ioutil.ReadFile(path)
		if err != nil {
			log.Printf("无法读取文件 %s: %v", path, err)
			return nil
		}

		fmt.Println("原始图片大小:", len(imgData), "字节")

		// 如果原始图片大小已经小于或等于目标大小，则直接复制图片
		if len(imgData) <= *maxSize {
			relativePath, err := filepath.Rel(*inputDir, path)
			if err != nil {
				log.Printf("无法获取相对路径 %s: %v", path, err)
				return nil
			}
			outputPath := filepath.Join(*outputDir, relativePath)

			// 确保输出目录存在
			outputDirPath := filepath.Dir(outputPath)
			if err := os.MkdirAll(outputDirPath, os.ModePerm); err != nil {
				log.Printf("无法创建输出目录 %s: %v", outputDirPath, err)
				return nil
			}

			// 保存图片到新目录，格式改为 .jpg
			outputPath = changeFileExtension(outputPath, ".jpg")
			if err := saveImage(outputPath, imgData); err != nil {
				log.Printf("无法保存图片 %s: %v", outputPath, err)
				return nil
			}

			fmt.Printf("图片 %s 已复制为 %s\n", path, outputPath)
			return nil
		}

		// 解码图片
		img, _, err := image.Decode(bytes.NewReader(imgData))
		if err != nil {
			log.Printf("无法解码图片 %s: %v", path, err)
			return nil
		}

		// 压缩图片到指定大小
		compressedImage, err := compressImageToSize(img, *maxSize)
		if err != nil {
			log.Printf("无法压缩图片 %s: %v", path, err)
			return nil
		}

		fmt.Println("压缩后图片大小:", len(compressedImage), "字节")

		// 构建输出文件路径
		relativePath, err := filepath.Rel(*inputDir, path)
		if err != nil {
			log.Printf("无法获取相对路径 %s: %v", path, err)
			return nil
		}
		outputPath := filepath.Join(*outputDir, relativePath)

		// 确保输出目录存在
		outputDirPath := filepath.Dir(outputPath)
		if err := os.MkdirAll(outputDirPath, os.ModePerm); err != nil {
			log.Printf("无法创建输出目录 %s: %v", outputDirPath, err)
			return nil
		}

		// 保存压缩后的图片，格式改为 .jpg
		outputPath = changeFileExtension(outputPath, ".jpg")
		if err := saveImage(outputPath, compressedImage); err != nil {
			log.Printf("无法保存图片 %s: %v", outputPath, err)
			return nil
		}

		fmt.Printf("图片 %s 已压缩并保存为 %s\n", path, outputPath)

		return nil
	})

	if err != nil {
		log.Fatalf("遍历文件夹时出错: %v", err)
	}

	fmt.Println("图片处理完成")
}

// compressImageToSize 压缩图片到指定大小
func compressImageToSize(img image.Image, maxSize int) ([]byte, error) {
	var buf bytes.Buffer
	quality := 100

	for {
		fmt.Println("准备压缩 quality:", quality)
		buf.Reset()

		// 使用 image/jpeg 进行质量调整和编码
		err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: quality})
		if err != nil {
			return nil, err
		}

		if buf.Len() <= maxSize {
			return buf.Bytes(), nil
		}

		if quality <= 30 {
			break
		}

		quality -= 2
	}

	// 如果仅通过质量压缩无法达成目标，削减四分之一的图片宽高
	bounds := img.Bounds()
	width := bounds.Dx()
	height := bounds.Dy()

	for {
		if width <= 0 || height <= 0 {
			return nil, fmt.Errorf("图片尺寸已达到最小值，无法进一步压缩")
		}

		fmt.Printf("削减图片尺寸: %dx%d\n", width, height)

		// 使用第三方库进行图片调整大小
		resizedImg := resizeImage(img, width, height)

		quality = 100
		for {
			fmt.Println("准备压缩 quality:", quality)
			buf.Reset()

			// 使用 image/jpeg 进行质量调整和编码
			err := jpeg.Encode(&buf, resizedImg, &jpeg.Options{Quality: quality})
			if err != nil {
				return nil, err
			}

			if buf.Len() <= maxSize {
				return buf.Bytes(), nil
			}

			if quality <= 20 {
				break
			}

			quality -= 5
		}

		// 继续削减图片尺寸
		width = width * 7 / 8
		height = height * 7 / 8

		// 确保新的尺寸不为零
		if width <= 0 {
			width = 1
		}
		if height <= 0 {
			height = 1
		}
	}
}

// resizeImage 调整图片大小
func resizeImage(img image.Image, width, height int) image.Image {
	bounds := img.Bounds()
	aspectRatio := float64(bounds.Dx()) / float64(bounds.Dy())

	if width == 0 {
		width = int(float64(height) * aspectRatio)
	}
	if height == 0 {
		height = int(float64(width) / aspectRatio)
	}

	newImg := image.NewRGBA(image.Rect(0, 0, width, height))
	draw.CatmullRom.Scale(newImg, newImg.Bounds(), img, img.Bounds(), draw.Over, nil)
	return newImg
}

// saveImage 保存图片到指定路径
func saveImage(path string, imgData []byte) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	_, err = file.Write(imgData)
	if err != nil {
		return err
	}

	return nil
}

// changeFileExtensionTxt 更改文件扩展名
func changeFileExtension(path, ext string) string {
	ext = strings.ToLower(ext)
	if ext[0] != '.' {
		ext = "." + ext
	}
	dir, file := filepath.Split(path)
	extOnly := filepath.Ext(file)
	file = file[:len(file)-len(extOnly)] + ext
	return filepath.Join(dir, file)
}
